import csv
from datetime import datetime, timedelta
import random

def generate_period_data(start_date, months=10):
    """
    生成带症状的经期数据
    """
    cycles = []
    current_date = datetime.strptime(start_date, '%Y-%m-%d')
    
    # 周期参数
    cycle_mean = 28
    cycle_std = 3
    period_mean = 5
    period_std = 1
    
    # 症状参数
    pain_weights = [30, 40, 20, 10]  # 无/轻微/中等/严重
    stress_weights = [40, 30, 30]    # 低/中/高
    
    for i in range(months * 2):
        # 生成周期数据
        cycle_length = max(21, min(35, int(random.gauss(cycle_mean, cycle_std))))
        period_days = max(3, min(7, int(random.gauss(period_mean, period_std))))
        
        # 生成每日症状数据
        symptoms = []
        for day in range(cycle_length):
            cycle_day = day % 28  # 强制周期为28天
            is_menstruating = 1 if cycle_day < period_days else 0
            date = current_date + timedelta(days=day)
            
            # 只在经期生成症状数据
            if is_menstruating:
                pain = random.choices(['无','轻微','中等','严重'], weights=pain_weights)[0]
                stress = random.choices(['低','中','高'], weights=stress_weights)[0]
            else:
                pain = '无'
                stress = random.choices(['低','中','高'], weights=[50,30,20])[0]
                
            symptoms.append({
                "date": date.strftime('%Y-%m-%d'),
                "is_menstruating": is_menstruating,
                "pain_level": pain,
                "stress_level": stress
            })
        
        # 保存周期数据
        cycles.append({
            "start_date": current_date.strftime('%Y-%m-%d'),
            "cycle_length": cycle_length,
            "period_days": period_days,
            "symptoms": symptoms
        })
        
        current_date += timedelta(days=cycle_length)
        if (current_date - datetime.strptime(start_date, '%Y-%m-%d')).days > months * 30:
            break

    # 写入完整周期数据
    with open('period_full_data.csv', 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        #date,is_period,pain,stress,start_date
        writer.writerow(['date', 'is_period', 'pain', 'stress', 'start_date'])
        for cycle in cycles:
            for day in cycle['symptoms']:
                writer.writerow([
                    day['date'],
                    day['is_menstruating'],
                    day['pain_level'],
                    day['stress_level'],
                    cycle['start_date']
                ])
    
    print(f"生成{len(cycles)}个周期的数据到period_full_data.csv")

# 使用示例
generate_period_data('2023-01-01', months=10) 